from sklearn.datasets.samples_generator import make_blobs
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans
from sklearn.mixture import GaussianMixture
import numpy as np

plt.rcParams['axes.unicode_minus'] = False  # 解决不显示负数问题

def datasets():
    X, _ = make_blobs(n_samples=800, centers=4, random_state=11)
    plt.scatter(X[:, 0], X[:, 1])
    plt.show()

    rng = np.random.RandomState(13)
    Y = np.dot(X, rng.randn(2, 2))
    plt.scatter(Y[:, 0], Y[:, 1])
    plt.show()

    return X, Y


def GmmKmean(dataX, dataY):
    """
    GMM算法与Kmeans算法对比
    :return:
    """
    kmeans = KMeans(n_clusters=4)
    kmeans.fit(dataX)
    y_kmeans = kmeans.predict(dataX)

    plt.scatter(dataX[:, 0], dataX[:, 1], c=y_kmeans)
    plt.show()
    centers = kmeans.cluster_centers_
    print(centers)

    gmm = GaussianMixture(n_components=4, random_state=1)
    gmm.fit(dataX)
    labels = gmm.predict(dataX)
    plt.scatter(dataX[:, 0], dataX[:, 1], c=labels)
    plt.show()

    kmeansy = KMeans(n_clusters=4, random_state=1)
    kmeansy.fit(dataY)
    datay_kmeans = kmeansy.predict(dataY)
    plt.scatter(dataY[:, 0], dataY[:, 1], c=datay_kmeans, s=40, cmap='viridis')
    plt.show()

    # gmmY = GaussianMixture(n_components=4)
    gmm.fit(dataY)
    labelsY = gmm.predict(dataY)
    plt.scatter(dataY[:, 0], dataY[:, 1], c=labelsY)
    plt.show()

    return None

if __name__ == "__main__":
    dataX, dataY = datasets()
    GmmKmean(dataX, dataY)
